package main

import (
	"encoding/json"
	"math"
	"math/rand"
	"os"
	"strings"
	"unicode"
)

// 辅助函数: 分割提示文本
func splitPrompt(prompt string) []string {
	var words []string
	runes := []rune(prompt)
	for i := 0; i < len(runes); {
		// 跳过空格
		if runes[i] == ' ' {
			i++
			continue
		}

		// 检查是否是中文字符
		if unicode.Is(unicode.Han, runes[i]) {
			// 中文字符按字分割
			words = append(words, string(runes[i]))
			i++
		} else {
			// 非中文字符按连续字符分割
			j := i
			for j < len(runes) && runes[j] != ' ' && !unicode.Is(unicode.Han, runes[j]) {
				j++
			}
			if j > i {
				words = append(words, string(runes[i:j]))
			}
			i = j
		}
	}
	return words
}

// 辅助函数: 判断是否是标点符号
func isPunctuation(word string) bool {
	punctuations := []string{"。", "，", "！", "？", ",", ".", "!", "?", ";", ":", "'", "\"", "《", "》", "（", "）", "【", "】", "{", "}", "[", "]", "(", ")"}
	for _, p := range punctuations {
		if word == p {
			return true
		}
	}
	return false
}

// 辅助函数: Tanh激活函数
func Tanh(x float32) float32 {
	return float32(math.Tanh(float64(x)))
}

// 辅助函数: Softmax函数
func Softmax(x []float32) []float32 {
	maxVal := x[0]
	for _, val := range x {
		if val > maxVal {
			maxVal = val
		}
	}

	expVals := make([]float32, len(x))
	sumExp := float32(0)
	for i, val := range x {
		expVals[i] = float32(math.Exp(float64(val - maxVal)))
		sumExp += expVals[i]
	}

	result := make([]float32, len(x))
	for i, val := range expVals {
		result[i] = val / sumExp
	}
	return result
}

// 辅助函数: 交叉熵损失
func CrossEntropyLoss(predictions []float32, targets []float32) float32 {
	loss := float32(0)
	for i, p := range predictions {
		if p > 0 {
			loss -= targets[i] * float32(math.Log(float64(p)))
		}
	}
	return loss
}

// 辅助函数: Dropout
func Dropout(x [][][]float32, rate float32) [][][]float32 {
	if rate <= 0 {
		return x
	}

	result := make([][][]float32, len(x))
	for b := 0; b < len(x); b++ {
		result[b] = make([][]float32, len(x[b]))
		for i := 0; i < len(x[b]); i++ {
			result[b][i] = make([]float32, len(x[b][i]))
			for j := 0; j < len(x[b][i]); j++ {
				if rand.Float32() > rate {
					result[b][i][j] = x[b][i][j] / (1 - rate)
				} else {
					result[b][i][j] = 0
				}
			}
		}
	}
	return result
}

// 辅助函数: 位置编码
func PositionalEncoding(maxLen, dModel int) [][]float32 {
	posEnc := make([][]float32, maxLen)
	for pos := 0; pos < maxLen; pos++ {
		posEnc[pos] = make([]float32, dModel)
		for i := 0; i < dModel; i += 2 {
			// 计算正弦和余弦位置编码
			angle := float64(pos) / math.Pow(10000, float64(i)/float64(dModel))
			posEnc[pos][i] = float32(math.Sin(angle))
			if i+1 < dModel {
				posEnc[pos][i+1] = float32(math.Cos(angle))
			}
		}
	}
	return posEnc
}

// 辅助函数: 检查早停条件
func CheckEarlyStopping(lossHistory []float32, patience int, minDelta float32) bool {
	if len(lossHistory) < patience+1 {
		return false
	}

	// 获取最近patience个损失值
	recentLosses := lossHistory[len(lossHistory)-patience-1:]

	// 检查是否连续patience个epoch损失都没有下降至少minDelta
	firstLoss := recentLosses[0]
	for i := 1; i <= patience; i++ {
		if firstLoss-recentLosses[i] > minDelta {
			return false
		}
	}

	return true
}

// 辅助函数: 加载ID映射表
func LoadIDMapping(path string) (map[string]int, error) {
	mapping := make(map[string]int)

	// 检查文件是否存在
	if _, err := os.Stat(path); os.IsNotExist(err) {
		// 文件不存在，返回空映射
		return mapping, nil
	}

	// 读取文件
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	// 解析JSON
	if err := json.Unmarshal(data, &mapping); err != nil {
		return nil, err
	}

	return mapping, nil
}

// 辅助函数: 保存ID映射表
func SaveIDMapping(mapping map[string]int, path string) error {
	// 转换为JSON
	data, err := json.MarshalIndent(mapping, "", "  ")
	if err != nil {
		return err
	}

	// 写入文件
	if err := os.WriteFile(path, data, 0644); err != nil {
		return err
	}

	return nil
}

// 辅助函数: 获取字符ID
// 采用万国码十进制+256位偏移量方案
func GetCharID(char string, mapping map[string]int) int {
	id, ok := mapping[char]
	if !ok {
		// 256位偏移量保护区内开始//
		// 数字 10*1
		if len(char) == 1 {
			if char >= "0" && char <= "9" {
				id = int(char[0] - '0')
				mapping[char] = id
				return id
			}
		}

		if char == "<unk>" {
			id = 10
			mapping[char] = id
			return id
		}
		if char == "<pad>" {
			id = 11
			mapping[char] = id
			return id
		}
		if char == "<bos>" {
			id = 12
			mapping[char] = id
			return id
		}
		if char == "<eos>" {
			id = 13
			mapping[char] = id
			return id
		}
		if char == "<mask>" {
			id = 14
			mapping[char] = id
			return id
		}
		if char == "<sep>" {
			id = 15
			mapping[char] = id
			return id
		}
		if char == "<cls>" {
			id = 16
			mapping[char] = id
			return id
		}

		// 英文字母 26个*2
		if len(char) == 1 {
			if char >= "a" && char <= "z" {
				id = int(char[0]-'a') + 26
				mapping[char] = id
				return id
			}
			if char >= "A" && char <= "Z" {
				id = int(char[0]-'A') + 52 // 从52开始编号
				mapping[char] = id
				return id
			}
		}
		//256位偏移量保护区内结束//

		// 对于单个字符，使用万国码+256偏移量
		if len(char) == 3 { // UTF8编码单个中文字符
			r := []rune(char)[0]
			id = int(r) + 256 // 256位偏移量
		} else {
			// 对于多字符词组，使用映射表长度+10万汉字单字预留+256+1作为ID
			id = len(mapping) + 100000 + 256 + 1 //临时解决方案

			// 对于多字符词组，使用字符的万国码拼接后哈希+256+1作为ID
			// var hash int64 = 0
			// for _, r := range char {
			// 	hash = hash*31 + int64(r)
			// }
			// // 限制ID在0-200000范围内
			// id = int(math.Abs(float64(hash)))%200000 + 256 + 1
		}
		mapping[char] = id
	}
	return id
}

// 辅助函数: 提取N-Gram词组
func ExtractNGrams(corpus []string, ngramSize int, minFrequency int) map[string]int {
	ngrams := make(map[string]int)

	for _, sentence := range corpus {
		words := splitPrompt(sentence)
		if len(words) < ngramSize {
			continue
		}

		for i := 0; i <= len(words)-ngramSize; i++ {
			ngram := strings.Join(words[i:i+ngramSize], " ")
			ngrams[ngram]++
		}
	}

	// 过滤低频N-Gram
	filteredNGrams := make(map[string]int)
	for ngram, freq := range ngrams {
		if freq >= minFrequency {
			filteredNGrams[ngram] = freq
		}
	}

	return filteredNGrams
}
